#!/usr/bin/env python
# coding: utf-8

# # 动态图模式应用
# 
# ## 概述
# 
# luojianet_ms的针对动态图和静态图两种模式在调试和运行方面做了不同的优化:
# 
# - 动态图模式：也称PyNative模式，将神经网络中的各个算子逐一下发执行，方便用户编写和调试神经网络模型。
# - 静态图模式：也称Graph模式或者图模式，将神经网络模型编译成一整张图，然后下发执行。该模式利用图优化等技术提高运行性能，同时有助于规模部署和跨平台运行。
# 
# 在动态图模式下，luojianet_ms支持执行单算子、普通函数和网络，以及单独求梯度的操作，下面我们将通过示例代码详细介绍这几种操作的使用方法和注意事项。
# 
# ## 动态图模式下的操作
# 
# 首先，我们导入相关依赖，并设置运行模式为动态图模式：

# In[1中低阶API实现深度学习]:


import numpy as np
import luojianet_ms.ops as ops
import luojianet_ms.nn as nn
from luojianet_ms import Tensor, context

context.set_context(mode=context.PYNATIVE_MODE)


# ### 执行单算子
# 
# 下面为执行加法算子[luojianet_ms.ops.Add](http://58.48.42.237/luojiaNet/luojiaNetapi/#luojianet_ms.ops.Add)的示例代码：

# In[2高级数据集管理]:


add = ops.Add()
x = Tensor(np.array([1, 2]).astype(np.float32))
y = Tensor(np.array([3, 5]).astype(np.float32))
z = add(x, y)
print("x:", x.asnumpy(), "\ny:", y.asnumpy(), "\nz:", z.asnumpy())


# ### 执行函数
# 
# 执行自定义函数`add_func`，示例代码如下：

# In[3图像处理]:


add = ops.Add()

def add_func(x, y):
    z = add(x, y)
    z = add(z, x)
    return z

x = Tensor(np.array([1, 2]).astype(np.float32))
y = Tensor(np.array([3, 5]).astype(np.float32))
z = add_func(x, y)
print("x:", x.asnumpy(), "\ny:", y.asnumpy(), "\nz:", z.asnumpy())


# ### 执行网络
# 
# 执行自定义网络`Net`，在forward中定义网络结构，示例代码如下：

# In[4自然语言]:


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.mul = ops.Mul()

    def forward(self, x, y):
        return self.mul(x, y)

net = Net()
x = Tensor(np.array([1.0, 2.0, 3.0]).astype(np.float32))
y = Tensor(np.array([4.0, 5.0, 6.0]).astype(np.float32))
z = net(x, y)

print("x:", x.asnumpy(), "\ny:", y.asnumpy(), "\nz:", z.asnumpy())


# ## 提升动态图模式性能
# 
# 为了提高动态图模式下的前向计算任务执行速度，luojianet_ms提供了`ms_function`装饰器，可以通过修饰Python函数或者Python类的成员函数使其被编译成计算图，通过图优化等技术提高运行速度。
# 
# ## 动态图模式下同步执行
# 
# 在动态图模式下，为了提升性能，算子在device上使用了异步执行方式，因此在算子执行错误的时候，错误信息可能会在程序执行到最后才显示。针对这种情况，luojianet_ms增加了一个pynative_synchronize的设置来控制算子device上是否使用异步执行。
# 
# 动态图模式下算子默认为异步执行，可以通过设置context来控制是否异步执行。当算子执行失败时，可以方便地通过调用栈看到出错的代码位置。示例代码如下：
# 
# ```python
# import numpy as np
# import luojianet_ms.context as context
# import luojianet_ms.nn as nn
# from luojianet_ms import Tensor
# from luojianet_ms import dtype as mstype
# import luojianet_ms.ops as ops
# 
# # 通过设置pynative_synchronize来使算子同步执行
# context.set_context(mode=context.PYNATIVE_MODE, pynative_synchronize=True)
# 
# class Net(nn.Module):
#     def __init__(self):
#         super(Net, self).__init__()
#         self.get_next = ops.GetNext([mstype.float32], [(1中低阶API实现深度学习, 1中低阶API实现深度学习)], 1中低阶API实现深度学习, "test")
# 
#     def forward(self, x1,):
#         x = self.get_next()
#         x = x + x1
#         return x
# 
# context.set_context()
# x1 = np.random.randn(1中低阶API实现深度学习, 1中低阶API实现深度学习).astype(np.float32)
# net = Net()
# output = net(Tensor(x1))
# print(output.asnumpy())
# ```
# 
# 输出：此时算子为同步执行，当算子执行错误时，可以看到完整的调用栈，找到出错的代码行。
# 
# ```text
# Traceback (most recent call last):
#   File "test.py", line 24, in <module>
#     output = net(Tensor(x1))
#   File ".../luojianet_ms/nn/cell.py", line 602, in __call__
#     raise err
#   File ".../luojianet_ms/nn/cell.py", line 599, in __call__
#     output = self._run_forward(cast_inputs, kwargs)
#   File ".../luojianet_ms/nn/cell.py", line 429, in _run_forward
#     output = self.forward(*cast_inputs, **kwargs)
#   File "test.py", line 17, in forward
#     x = self.get_next()
#   File ".../luojianet_ms/ops/primitive.py", line 294, in __call__
#     return _run_op(self, self.name, args)
#   File ".../luojianet_ms/common/api.py", line 90, in wrapper
#     results = fn(*arg, **kwargs)
#   File ".../luojianet_ms/ops/primitive.py", line 754, in _run_op
#     output = real_run_op(obj, op_name, args)
# RuntimeError: luojianet_ms/ccsrc/plugin/device/gpu/kernel/data/dataset_iterator_kernel.cc:139 Launch] For 'GetNext', gpu Queue(test) Open Failed: 2高级数据集管理
# ```
# 
# ## Hook功能
# 
# 调试深度学习网络是每一个深度学习领域的从业者需要面对且投入精力较大的工作。由于深度学习网络隐藏了中间层算子的输入、输出数据以及反向梯度，只提供网络输入数据（特征量、权重）的梯度，导致无法准确地感知中间层算子的数据变化，从而降低了调试效率。为了方便用户准确、快速地对深度学习网络进行调试，luojianet_ms在动态图模式下设计了Hook功能，**使用Hook功能可以捕获中间层算子的输入、输出数据以及反向梯度**。
# 
# 目前，动态图模式下提供了四种形式的Hook功能，分别是：HookBackward算子和在Cell对象上进行注册的register_forward_pre_hook、register_forward_hook、register_backward_hook功能。
# 
# ### HookBackward算子
# 
# HookBackward将Hook功能以算子的形式实现。用户初始化一个HookBackward算子，将其安插到深度学习网络中需要捕获梯度的位置。在网络正向执行时，HookBackward算子将输入数据不做任何修改后原样输出；在网络反向传播梯度时，在HookBackward上注册的Hook函数将会捕获反向传播至此的梯度。用户可以在Hook函数中自定义对梯度的操作，比如打印梯度，或者返回新的梯度。
# 
# 示例代码:

# In[5]:


import luojianet_ms
from luojianet_ms import ops
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def hook_fn(grad_out):
    """打印梯度"""
    print("hook_fn print grad_out:", grad_out)

grad_all = GradOperation(get_all=True)
hook = ops.HookBackward(hook_fn)
def hook_test(x, y):
    z = x * y
    z = hook(z)
    z = z * y
    return z

def net(x, y):
    return grad_all(hook_test)(x, y)

output = net(Tensor(1, luojianet_ms.float32), Tensor(2, luojianet_ms.float32))
print("output:", output)


# 更多HookBackward算子的说明可以参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/)。
# 
# ### Cell对象的register_forward_pre_hook功能
# 
# 用户可以在Cell对象上使用`register_forward_pre_hook`函数来注册一个自定义的Hook函数，用来捕获正向传入该Cell对象的数据。该功能在静态图模式下和在使用`ms_function`修饰的Cell对象上不起作用。`register_forward_pre_hook`函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的`handle`对象。用户可以通过调用`handle`对象的`remove()`函数来删除与之对应的Hook函数。每一次调用`register_forward_pre_hook`函数，都会返回一个不同的`handle`对象。Hook函数应该按照以下的方式进行定义。

# In[6]:


def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)


# 这里的cell_id是Cell对象的名称以及ID信息，inputs是正向传入到Cell对象的数据。因此，用户可以使用register_forward_pre_hook函数来捕获网络中某一个Cell对象的正向输入数据。用户可以在Hook函数中自定义对输入数据的操作，比如查看、打印数据，或者返回新的输入数据给当前的Cell对象。如果在Hook函数中对Cell对象的原始输入数据进行计算操作后，再作为新的输入数据返回，这些新增的计算操作将会同时作用于梯度的反向传播。
# 
# 示例代码:

# In[7]:


import numpy as np
import luojianet_ms
import luojianet_ms.nn as nn
import luojianet_ms.ops as ops
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    input_x = inputs[0]
    return input_x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)

    def forward(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

grad = GradOperation(get_all=True)
net = Net()

x = Tensor(np.ones([1]).astype(np.float32))
y = Tensor(np.ones([1]).astype(np.float32))

output = net(x, y)
print(output)
gradient = grad(net)(x, y)
print(gradient)
net.handle.remove()
gradient = grad(net)(x, y)
print(gradient)


# 用户如果在Hook函数中直接返回新创建的数据，而不是返回由原始输入数据经过计算后得到的数据，那么梯度的反向传播将会在该Cell对象上截止。
# 
# 示例代码:

# In[8]:


import numpy as np
import luojianet_ms
import luojianet_ms.nn as nn
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    return Tensor(np.ones([1]).astype(np.float32))

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)

    def forward(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

grad = GradOperation(get_all=True)
net = Net()

x = Tensor(np.ones([1]).astype(np.float32))
y = Tensor(np.ones([1]).astype(np.float32))

gradient = grad(net)(x, y)
print(gradient)


# 为了避免脚本在切换到图模式时运行失败，不建议在Cell对象的 `forward` 函数中调用 `register_forward_pre_hook` 函数和 `handle` 对象的 `remove()` 函数。在动态图模式下，如果在Cell对象的 `forward` 函数中调用 `register_forward_pre_hook` 函数，那么Cell对象每次运行都将新注册一个Hook函数。
# 
# 更多关于Cell对象的 `register_forward_pre_hook` 功能的说明可以参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/#luojianet_ms.nn.Module.register_forward_pre_hook)。
# 
# ### Cell对象的register_forward_hook功能
# 
# 用户可以在Cell对象上使用`register_forward_hook`函数来注册一个自定义的Hook函数，用来捕获正向传入Cell对象的数据和Cell对象的输出数据。该功能在静态图模式下和在使用`ms_function`修饰的Cell对象上不起作用。`register_forward_hook`函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的`handle`对象。用户可以通过调用`handle`对象的`remove()`函数来删除与之对应的Hook函数。每一次调用`register_forward_hook`函数，都会返回一个不同的`handle`对象。Hook函数应该按照以下的方式进行定义。
# 
# 示例代码:

# In[9]:


def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)


# 这里的`cell_id`是Cell对象的名称以及ID信息，`inputs`是正向传入到Cell对象的数据，`outputs`是Cell对象的正向输出数据。因此，用户可以使用`register_forward_hook`函数来捕获网络中某一个Cell对象的正向输入数据和输出数据。用户可以在Hook函数中自定义对输入、输出数据的操作，比如查看、打印数据，或者返回新的输出数据。如果在Hook函数中对Cell对象的原始输出数据进行计算操作后，再作为新的输出数据返回，这些新增的计算操作将会同时作用于梯度的反向传播。
# 
# 示例代码:

# In[10]:


import numpy as np
import luojianet_ms
import luojianet_ms.nn as nn
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)
    outputs = outputs + outputs
    return outputs

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_hook(forward_hook_fn)

    def forward(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

grad = GradOperation(get_all=True)
net = Net()

x = Tensor(np.ones([1]).astype(np.float32))
y = Tensor(np.ones([1]).astype(np.float32))

gradient = grad(net)(x, y)
print(gradient)
net.handle.remove()
gradient = grad(net)(x, y)
print(gradient)


# 用户如果在Hook函数中直接返回新创建的数据，而不是将原始的输出数据经过计算后，将得到的新输出数据返回，那么梯度的反向传播将会在该Cell对象上截止。该现象可以参考`register_forward_pre_hook`函数的用例说明。
# 为了避免脚本在切换到图模式时运行失败，不建议在Cell对象的`forward`函数中调用`register_forward_hook`函数和`handle`对象的`remove()`函数。在动态图模式下，如果在Cell对象的`forward`函数中调用`register_forward_hook`函数，那么Cell对象每次运行都将新注册一个Hook函数。
# 
# 更多关于Cell对象的`register_forward_hook`功能的说明可以参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/#luojianet_ms.nn.Module.register_forward_hook)。
# 
# ### Cell对象的register_backward_hook功能
# 
# 用户可以在Cell对象上使用`register_backward_hook`函数来注册一个自定义的Hook函数，用来捕获网络反向传播时与Cell对象相关联的梯度。该功能在图模式下或者在使用`ms_function`修饰的Cell对象上不起作用。`register_backward_hook`函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的`handle`对象。用户可以通过调用`handle`对象的`remove()`函数来删除与之对应的Hook函数。每一次调用`register_backward_hook`函数，都会返回一个不同的`handle`对象。
# 
# 与HookBackward算子所使用的自定义Hook函数有所不同，`register_backward_hook`使用的Hook函数的入参中，包含了表示Cell对象名称与id信息的`cell_id`、反向传入到Cell对象的梯度、以及Cell对象的反向输出的梯度。
# 
# 示例代码:

# In[11]:


def backward_hook_function(cell_id, grad_input, grad_output):
    print(grad_input)
    print(grad_output)


# 这里的`cell_id`是Cell对象的名称以及ID信息，`grad_input`是网络反向传播时，传入到Cell对象的梯度，它对应于正向过程中下一个算子的反向输出梯度；`grad_output`是Cell对象反向输出的梯度。因此，用户可以使用`register_backward_hook`函数来捕获网络中某一个Cell对象的反向传入和反向输出梯度。用户可以在Hook函数中自定义对梯度的操作，比如查看、打印梯度，或者返回新的输出梯度。如果需要在Hook函数中返回新的输出梯度时，返回值必须是`tuple`的形式。
# 
# 示例代码:

# In[12]:


import numpy as np
import luojianet_ms
import luojianet_ms.nn as nn
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def backward_hook_function(cell_id, grad_input, grad_output):
    print(grad_input)
    print(grad_output)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 2, kernel_size=2, stride=1, padding=0, weight_init="ones", pad_mode="valid")
        self.bn = nn.BatchNorm2d(2, momentum=0.99, eps=0.00001, gamma_init="ones")
        self.handle = self.bn.register_backward_hook(backward_hook_function)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

net = Net()
grad_all = GradOperation(get_all=True)
output = grad_all(net)(Tensor(np.ones([1, 1, 2, 2]).astype(np.float32)))
print(output)
net.handle.remove()
output = grad_all(net)(Tensor(np.ones([1, 1, 2, 2]).astype(np.float32)))
print("-------------\n", output)


# 当 `register_backward_hook` 函数和 `register_forward_pre_hook` 函数、 `register_forward_hook` 函数同时作用于同一Cell对象时，如果 `register_forward_pre_hook` 和 `register_forward_hook` 函数中有添加其他算子进行数据处理，这些新增算子会在Cell对象执行前或者执行后参与数据的正向计算，但是这些新增算子的反向梯度不在 `register_backward_hook` 函数的捕获范围内。 `register_backward_hook` 中注册的Hook函数仅捕获原始Cell对象的输入、输出梯度。
# 
# 示例代码:

# In[13]:


import numpy as np
import luojianet_ms
import luojianet_ms.nn as nn
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    input_x = inputs[0]
    return input_x

def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)
    outputs = outputs + outputs
    return outputs

def backward_hook_fn(cell_id, grad_input, grad_output):
    print("grad input: ", grad_input)
    print("grad output: ", grad_output)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)
        self.handle2 = self.relu.register_forward_hook(forward_hook_fn)
        self.handle3 = self.relu.register_backward_hook(backward_hook_fn)

    def forward(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

net = Net()
grad = GradOperation(get_all=True)
gradient = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32)))
print(gradient)


# 这里的 `grad_input` 是梯度反向传播时传入`self.relu`的梯度，而不是传入 `forward_hook_fn` 函数中，新增的 `Add` 算子的梯度。这里的 `grad_output` 是梯度反向传播时 `self.relu` 反向输出的梯度，而不是 `forward_pre_hook_fn` 函数中新增 `Add` 算子的反向输出梯度。 `register_forward_pre_hook` 函数和 `register_forward_hook` 函数是在Cell对象执行前后起作用，不会影响Cell对象上反向Hook函数的梯度捕获范围。
# 为了避免脚本在切换到图模式时运行失败，不建议在Cell对象的 `forward` 函数中调用 `register_backward_hook` 函数和 `handle` 对象的 `remove()` 函数。在PyNative模式下，如果在Cell对象的 `forward` 函数中调用 `register_backward_hook` 函数，那么Cell对象每次运行都将新注册一个Hook函数。
# 
# 更多关于Cell对象的 `register_backward_hook` 功能的说明可以参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/#luojianet_ms.nn.Module.register_backward_hook)。
# 
# ## 自定义bprop功能
# 
# 用户可以自定义nn.Module对象的反向传播（计算）函数，从而控制nn.Module对象梯度计算的过程，定位梯度问题。自定义bprop函数的使用方法是：在定义的nn.Module对象里面增加一个用户自定义的bprop函数。训练的过程中会使用用户自定义的bprop函数来生成反向图。
# 
# 示例代码:

# In[14]:


import luojianet_ms
import luojianet_ms.nn as nn
from luojianet_ms import Tensor
from luojianet_ms import context
from luojianet_ms.ops import GradOperation

context.set_context(mode=context.PYNATIVE_MODE)

class Net(nn.Module):
    def forward(self, x, y):
        z = x * y
        z = z * y
        return z

    def bprop(self, x, y, out, dout):
        x_dout = x + y
        y_dout = x * y
        return x_dout, y_dout

grad_all = GradOperation(get_all=True)
output = grad_all(Net())(Tensor(1, luojianet_ms.float32), Tensor(2, luojianet_ms.float32))
print(output)

